#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
板块信息表模型 (Block)
"""

from sqlalchemy import Column, Integer, String, Index
from sqlalchemy.orm import relationship
from typing import List

from .BaseModel import BaseModel

class Block(BaseModel):
    """板块信息表模型"""
    __tablename__ = 'block'
    
    # 列名映射（CSV列名 -> 数据库字段名）
    aliases = {
        'blockCode': 'block_code',
        'blockName': 'block_name',
        'category': 'category',
        '#板块编号': 'block_code',
        '板块名称': 'block_name',
        '类别': 'category'
    }
    
    block_code = Column(String(20), primary_key=True)
    block_name = Column(String(100))
    category = Column(Integer, index=True)
    
    # 关联关系
    block_stocks = relationship("BlockInfo", back_populates="block")
    
    # 索引
    __table_args__ = (
        Index('idx_block_code', 'block_code'),
        Index('idx_block_category', 'category'),
    )

    @classmethod
    def get_session(cls):
        """获取数据库会话"""
        return cls._db_manager.get_stock_data_session()
    
    @classmethod
    def _get_existing_record(cls, session, data):
        """获取现有记录"""
        return session.query(cls).filter(cls.block_code == data['block_code']).first()
    
    @classmethod
    def _update_existing_record(cls, existing, data):
        """更新现有记录"""
        existing.block_name = data.get('block_name', existing.block_name)
        existing.category = data.get('category', existing.category)
    
    @classmethod
    def _to_dict(cls, result):
        """将结果转换为字典"""
        return {
            'block_code': result.block_code,
            'block_name': result.block_name,
            'category': result.category
        }
    
    @classmethod
    def get_all(cls) -> List['Block']:
        """获取所有板块数据"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            results = session.query(cls).all()
            session.close()
            return results
        except Exception as e:
            # 获取板块数据失败: {str(e)}
            session.close()
            return []
    
    @classmethod
    def import_from_dataframe(cls, df):
        """从DataFrame导入板块数据"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            success_count = 0
            
            for _, row in df.iterrows():
                try:
                    # 检查是否已存在相同记录
                    existing = session.query(cls).filter(
                        cls.block_code == row['block_code']
                    ).first()
                    
                    if existing:
                        # 更新现有记录
                        existing.block_name = row.get('block_name', existing.block_name)
                        existing.category = row.get('category', existing.category)
                    else:
                        # 创建新记录
                        block = cls(
                            block_code=row['block_code'],
                            block_name=row.get('block_name', ''),
                            category=row.get('category', 0)
                        )
                        session.add(block)
                    
                    success_count += 1
                    
                except Exception as e:
                    # 处理板块数据失败: {str(e)}
                    continue
            
            session.commit()
            session.close()
            return success_count
            
        except Exception as e:
            # 导入板块数据失败: {str(e)}
            session.rollback()
            session.close()
            return 0
    
    @classmethod
    def to_dataframe(cls):
        """转换为DataFrame"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            results = session.query(cls).all()
            
            data_list = []
            for result in results:
                data_list.append({
                    'block_code': result.block_code,
                    'block_name': result.block_name,
                    'category': result.category
                })
            
            session.close()
            
            import pandas as pd
            df = pd.DataFrame(data_list)
            # 从数据库加载了 {len(df)} 条板块数据
            return df
            
        except Exception as e:
            # 转换为DataFrame失败: {str(e)}
            session.close()
            import pandas as pd
            return pd.DataFrame()
    
    @classmethod
    def load_from_db(cls, db_path: str = "stock_data.db"):
        """从数据库加载板块数据（兼容方法）"""
        # 创建临时数据库管理器
        from .BaseModel import DatabaseManager
        db_manager = DatabaseManager(stock_data_db_path=db_path)
        cls.set_db_manager(db_manager)
        
        try:
            df = cls.to_dataframe()
            return df
        finally:
            db_manager.close_all_sessions()
